#!/usr/bin/perl

=head1 Name

 phylo_tree.pl - Using TreeBeST or PhyML or MrBayes to 
 reconstruct phylogenetic tree

=head1 Version

 Author: hejing@genomics.org.cn
 Version 1.4, 2010-09-24

=head1 Usage
 
 perl phylo_tree.pl <in.phy> [options]
  <in.phy>         input phylip format(sequential) sequence file;
  -type <str>      method: nj(default); ml; bayes;
  -outdir <str>    save results to this directory, default ./;
  -clean           delete temporary files, default not;
  -help            show help information;

 
 TreeBeST's relevant parameters(nj):
  -d <str>         data type: nt, nucleotide(default); aa, amino acid;
  -t <str>         model for codon & nt: ntmm; dn; ds; dm(default); 
                   model for aa: mm(default); jtt; kimura;
  -b <num>         bootstrap times, default 100; 


 PhyML's relevant parameters(ml):
  -d <str>         data type: nt, nucleotide(default); aa, amino acid;
  -m <str>         models for nt: HKY85(default); JC69; K80; F81; F84; TN93; GTR;
                   models for aa: WAG(default); JTT; MtREV; Dayhoff;
  -b <num>         set bootstrap(int>0) or aLRT method(-1,-2,-3,-4), default -4;
  -rates <str>     rates across sites: gamma(default); invgamma; equal; propinv;


 MrBayes's relevant parameters(bayes):
  -data <str>      data type: dna(default); protein;
  -nst <num>       nucleotide models: 6, GTR(default); 2, HKY; 1, JC;
  -aamodel <str>   amino-acid models: Poisson; Jones(default); Dayhoff; WAG; BLOSUM;
  -rates <str>     rates across sites: gamma(default); invgamma; equal; propinv;
  -ngen <num>      generation number, default 100000;
  -samplefreq <num>    sample one time for every 100 (default) generations;

=head1 Example 

 perl phylo_tree.pl <in.phy> -type nj -d nt
 perl phylo_tree.pl <in.phy> -type ml -d nt
 perl phylo_tree.pl <in.phy> -type bayes -data dna

=cut

use strict;
use Getopt::Long;
use FindBin qw($Bin);
use lib "$Bin";
use lib "$Bin/../lib";
use File::Basename;

my ($type, $outdir) = ("nj", "./");
my ($help, $clean);
my ($aamodel, $data, $rates, $nst, $ngen, $samplefreq);
my ($d, $m, $b, $t);
GetOptions(
    'type:s'   => \$type,
    'outdir:s' => \$outdir,
    'clean'    => \$clean,
    'help'     => \$help,

    'data:s'       => \$data,
    'aamodel:s'    => \$aamodel,
    'nst:i'        => \$nst,
    'rates:s'      => \$rates,
    'ngen:i'       => \$ngen,
    'samplefreq:i' => \$samplefreq,

    'd:s' => \$d,
    't:s' => \$t,
    'b:i' => \$b,

    'm:s' => \$m,
);
die `pod2text $0` if (@ARGV < 1 || $help);
my $file_phy = shift;
mkdir $outdir if (!-e $outdir);

if ($type =~ /bayes/) {
    $data       ||= "dna";
    $rates      ||= "gamma";
    $nst        ||= 6;
    $ngen       ||= 100000;
    $samplefreq ||= 100;
    $aamodel    ||= "Jones";

    my $burnin = ($ngen / $samplefreq) * 0.25;

    open(IN,  "$file_phy")          or die "Can't open :$file_phy\n";
    open(OUT, ">$outdir/input.nex") or die "Can't open :$outdir/input.nex\n";
    my $firstline = <IN>;
    my ($species_number, $sequence_length);
    if ($firstline =~ /^\s*(\d+)\s+(\d+)/) {
        $species_number = $1, $sequence_length = $2;
        print OUT
"#NEXUS\n\nbegin data;\n\tdimensions ntax =$species_number nchar=$sequence_length;";
    }
    if ($sequence_length < 15000) {
        print OUT
"\n\tformat datatype=$data interleave =no missing=? gap=-;\n\tmatrix\n";
        while (my $eachline = <IN>) {
            if ($eachline =~ /^\w+\s+[\w-]/) {
                print OUT "\t$eachline";
            }
        }
    }
    else {
        my (@phylip, @name, @sequence, @aim, $change);
        @phylip = <IN>;
        foreach (@phylip) {
            my @temp = split /\s+/, $_;
            push @name,     $temp[0];
            push @sequence, $temp[1];
        }
        for (my $i = 0 ; $i < $sequence_length ; $i += 15000) {
            for (my $k = 0 ; $k < scalar(@sequence) ; $k++) {
                chomp $sequence[$k];
                my $change = substr($sequence[$k], $i, 15000);
                push @aim, "\t$name[$k]\t$change\n";
            }
        }
        print OUT
"\n\tformat datatype=$data interleave =yes missing=? gap=-;\n\tmatrix\n";
        foreach (@aim) {
            chomp;
            print OUT "$_\n";
        }
        close IN;
    }
    print OUT "\t;\nend;\n";
    print OUT "begin mrbayes;\n";
    print OUT "\tset autoclose=yes nowarn=yes;\n";
    if (lc($data) eq "dna") {
        print OUT "\tlset nst=$nst rates=$rates;\n";
    }
    elsif (lc($data) eq "protein") {
        print OUT "\tprset aamodelpr=fixed($aamodel);\n\tlset rates=$rates;\n";
    }
    print OUT
"\tmcmc ngen=$ngen printfreq=100 samplefreq=$samplefreq nchains=4 savebrlens=yes;\n";
    print OUT
      "\tsumt filename=$outdir/input.nex contype=allcompat burnin=$burnin;\n";
    print OUT "\tquit;\nend;\n";
    close OUT;
    system "$Bin/mb  -i $outdir/input.nex ";
    open(IN, "$outdir/input.nex.con")
      or die "Can't open :$outdir/input.nex.con\n";
    open(OUT, ">$outdir/tree.tre")
      or die "Can't open :$outdir/tree.tre\n";
    my $treestr;
    while (<IN>) {
        chomp;
        if ($_ =~ /^\s.+=/) {
            my @arr = split /=/, $_;
            $treestr = $arr[1];
            last;
        }
    }
    close IN;
    print OUT"$treestr";
    close OUT;
    ## To get B value ##
    open(OUT, ">$outdir/tree.nhx") or die "Can't open :$outdir/tree.nhx\n";
    my $str = $treestr;
    $str =~ s/\)([\d|.]+)\:([\d|.]+)/res($1, $2)/ge;
    print OUT $str;
    close OUT;

    sub res {
        my ($a, $b) = @_;
        my $c = $a * 100;
        #return "\)\:$b\[&&NHX:B=$c\]";
        return "\)\:$b";
    }

    `rm  $outdir/tree.newick $outdir/input.nex*` if ($clean);

}
elsif ($type =~ /nj/) {

    $d ||= "nt";
    if    ($d eq "nt") { $t ||= "dm"; }
    elsif ($d eq "aa") { $t ||= "mm"; }
    $b ||= "100";
    open(IN, "$file_phy") or die "Can't open :$file_phy\n";
    <IN>;
    open(OUT, ">$outdir/in.mfa") or die "can't open :$outdir/in.mfa\n";
    while (<IN>) {
        chomp;
        my @arr1 = split /\s+/;
        print OUT ">$arr1[0]\n$arr1[1]\n";
    }
    close IN;
    close OUT;
	print "datatype = $d\nmodel = $t\nbootstrap = $b\n";

    ###call njtree##
    system
"/mnt/sdb/zengl/bin/rna-pipe/Compare-rna/ploy/treebest-1.9.2/treebest nj -t $t $outdir/in.mfa -b $b 1> $outdir/tree1.nhx 2> /dev/null";

    open(IN,  "$outdir/tree1.nhx") or die "Can't open :$outdir/tree1.nhx\n";
    open(OUT, ">$outdir/tree.nhx") or die "can't open :$outdir/tree.nhx\n";
    while (<IN>) {
        chomp;
        print OUT "$_";
    }
    close IN;
    close OUT;
    `rm $outdir/tree.newick $outdir/in.mfa` if ($clean);

}
elsif ($type =~ /ml/) {
    $d     ||= "nt";
    $b     ||= -4;
    $rates ||= "gamma";
    `cp $file_phy $outdir/temp.phy`; 
	if ($d eq "nt") {
        $m ||= "HKY85";
        my $same =
          "$Bin/phyml_3.0_linux -i $outdir/temp.phy -d $d -b $b -m $m -q";
        if    ($rates =~ /equal/)    { system "$same -t e"; }
        elsif ($rates =~ /gamma/)    { system "$same -a e -t e"; }
        elsif ($rates =~ /invgamma/) { system "$same -a e -c 4 -v e -t e"; }
        elsif ($rates =~ /propinv/)  { system "$same -v e -t e"; }
    }
    elsif ($d eq "aa") {
        $m ||= "WAG";
        my $same =
          "$Bin/phyml_3.0_linux -i $outdir/temp.phy -d $d -b $b -m $m -q";
        if    ($rates =~ /equal/)    { system "$same"; }
        elsif ($rates =~ /gamma/)    { system "$same -a e"; }
        elsif ($rates =~ /invgamma/) { system "$same -a e -c 4 -v e"; }
        elsif ($rates =~ /propinv/)  { system "$same -v e"; }

    }
    open(IN, "$outdir/temp.phy_phyml_tree.txt")
      or die "Can't open :$outdir/temp.phy_phyml_tree.txt\n";
    open(OUT, ">$outdir/tree.newick")
      or die "Can't open :$outdir/tree.newick\n";
    my $treestr = <IN>;
    print OUT "$treestr";
    close OUT;

    ##To get B value##
    open(OUT, ">$outdir/tree.nhx") or die "Can't open :$outdir/tree.nhx\n";
    chomp $treestr;
    if ($b > 0) {
        #$treestr =~ s/\)(\d+)\:([\d|.]+)/\)\:$2\[&&NHX:B=$1\]/g;
        $treestr =~ s/\)(\d+)\:([\d|.]+)/\)\:$2/g;
        print OUT"$treestr\n";
    }
    else {
        $treestr =~ s/\)([\d|.]+)\:([\d|.]+)/res($1, $2)/ge;
        print OUT "$treestr";

        sub res {
            my ($a, $b) = @_;
            my $c = $a * 100;
            #return "\)\:$b\[&&NHX:B=$c\]";
            return "\)\:$b";
        }
    }
    close IN;
    close OUT;
    `rm $outdir/tree.newick $outdir/temp.phy_phyml_*` if ($clean);
}
system "perl $Bin/draw_tree.pl $outdir/tree.nhx > $outdir/tree.svg";
